from pymongo import MongoClient
from pymongo.database import Database
import os
from typing import Optional
import redis
from redis import Redis

# MongoDB连接
def get_mongo_db() -> Database:
    """获取MongoDB数据库连接"""
    mongo_uri = os.getenv("MONGODB_URI")
    client = MongoClient(mongo_uri)
    db = client.get_database()
    return db

# Redis连接
def get_redis_client() -> Redis:
    """获取Redis客户端连接"""
    redis_host = os.getenv("REDIS_HOST")
    redis_port = int(os.getenv("REDIS_PORT", 6379))
    redis_password = os.getenv("REDIS_PASSWORD")
    redis_db = int(os.getenv("REDIS_DB", 0))
    
    client = redis.Redis(
        host=redis_host,
        port=redis_port,
        password=redis_password,
        db=redis_db,
        decode_responses=True
    )
    return client

# 数据库连接实例
mongo_db: Optional[Database] = None
redis_client: Optional[Redis] = None

def init_db():
    """初始化数据库连接"""
    global mongo_db, redis_client
    mongo_db = get_mongo_db()
    redis_client = get_redis_client()

# 关闭数据库连接
def close_db():
    """关闭数据库连接"""
    global mongo_db, redis_client
    if mongo_db:
        mongo_db.client.close()
    if redis_client:
        redis_client.close()
